import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image, ImageDraw, ImageFont

DIGITS_LETTERS = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
FONT_PATHS = [
    'C:/Windows/Fonts/arialbi.ttf',  # 粗斜体
    'C:/Windows/Fonts/arialbd.ttf',  # 粗体
    'C:/Windows/Fonts/ariali.ttf',  # 斜体
    # 'C:/Windows/Fonts/consola.ttf',
    # 'C:/Windows/Fonts/consolai.ttf',
    # 'C:/Windows/Fonts/consolai.ttf'
]


class VerCodeData(Dataset):
    def __init__(self, data_len=36 * 2000, noise_rate=1, out_size=(32, 32)):
        self.data_len = data_len
        self.noise_rate = noise_rate
        self.out_size = out_size

    def __getitem__(self, item):
        # 字符的大小，位置，字符
        r_fonttype = np.random.randint(0, len(FONT_PATHS))
        r_fontsize = np.random.randint(18, 32)
        r_px = np.random.randint(-6, 3)
        r_py = np.random.randint(-6, 3)
        r_dl_idx = np.random.randint(0, 36)
        r_dl = DIGITS_LETTERS[r_dl_idx]

        # 创建空白图，底色为白色，并向其写入字符
        fnt = ImageFont.truetype(FONT_PATHS[r_fonttype], r_fontsize)
        txt_size = fnt.getsize(r_dl)
        txt = Image.new('RGB', txt_size, (255, 255, 255))
        d = ImageDraw.Draw(txt)
        d.text((r_px, r_py), r_dl, font=fnt, fill=tuple(np.random.randint(0, 100, 3)))

        # 添加随机噪声，比率一般为1/10
        mask = np.random.randint(0, 10, (txt_size[1], txt_size[0], 3))
        mask = np.where(mask < self.noise_rate, True, False)
        img_noise = mask * np.random.randint(0, 100, (txt_size[1], txt_size[0], 3), dtype=np.uint8)
        # 将数据下采样，再上采样，以降低分辨率，保证与原始的验证码差不多
        dl_img = Image.fromarray(np.array(txt) - img_noise).resize(self.out_size)

        # 返回字符图片的数组，和它的标签
        return np.array(dl_img), r_dl_idx

    def __len__(self):
        return self.data_len
